def <- simstudy::defData(varname = "C", formula = 0.4, dist = "binary")
def <- simstudy::defData(def, "X", formula = "0.3 + 0.4 * C", dist = "binary")
def <- simstudy::defData(def, "e", formula = 0, variance = 2, dist = "normal")
def <- simstudy::defData(def, "Y0", formula = "2 * C + e", dist="nonrandom")
def <- simstudy::defData(def, "Y1", formula = "1 + 2 * C + e", dist="nonrandom")
def <- simstudy::defData(def, "Y_obs", formula = "Y0 + (Y1 - Y0) * X", dist = "nonrandom") #  = Y1*X + (1-X)*Y0

set.seed(2017)
dt <- genData(10000, def)

paste0("Calculated mean causal effect = ", mean(dt[, Y1] - dt[, Y0])) #mean difference of counterfactual outcomes

paste0("Calculated observed difference  = ", round(dt[X == 1, mean(Y_obs)] - dt[X == 0, mean(Y_obs)],2))

lm1 <- lm(Y_obs ~ X, data = dt)
tidy(lm1)

# df DAG 
dag_collid <- dagify(
  Y ~ X + e,
  Z ~ X + Y,
  exposure = "X",
  outcome = "Y",
  #latent = "e",
  # labels
  labels = c(
    "Z" = "Collider",
    "X" = "Exposure",
    "Y" = "Outcome",
    "e" = "Unobserved \nerror"),
  # positions
  coords = list(
    x = c(X = 0, Y= 4, Z = 2 , e = 5),
    y = c(X = 0, Y = 0, Z = 2, e = 1) 
  )) %>% 
  tidy_dagitty() 

# Plot DAG 
dag_col_p <-  dag_collid %>% 
  ggdag(., layout = "circle") +  # decided in dagify
  theme_dag_blank(plot.caption = element_text(hjust = 1)) +
  # Nodes 
  geom_dag_node(aes(color = name)) +  
  # Labels dodged to avoid overlap
  geom_dag_label(aes(label = label), color = "#4c4c4c", nudge_x = 0.7, nudge_y = 0.2) +  
  
  geom_dag_text(color="white") +
  labs(title = "Causal map with COLLIDER (Z)" , 
       #subtitle = "X = exposure, Y = outcome, Z = collider, e = random error", 
       caption = ) +
  # Map colors to specific nodes
  scale_color_manual(values = c("X" = "#9b6723", "Y" = "#72aed8", "Z" = "#d02e4c",
                                "e" = "#A6A6A6"), 
                     guide = "none")    # Remove legend

dag_col_p

library(dagitty) # Graphical Analysis of Structural Causal Models # Graphical Analysis of Structural Causal Models # Graphical Analysis of Structural Causal Models 
library(ggdag) # Analyze and Create Elegant Directed Acyclic Graphs # Analyze and Create Elegant Directed Acyclic Graphs # Analyze and Create Elegant Directed Acyclic Graphs 

# df DAG 
dag_confounder <- dagify(
  Y ~ X + Z + e,
  X ~  Z,
  exposure = "X",
  outcome = "Y",
  #latent = "e",
  # labels
  labels = c(
    "Z" = "Confounder",
    "X" = "Exposure",
    "Y" = "Outcome",
    "e" = "Unobserved \nerror"),
  # positions
  coords = list(
    x = c(X = 0, Y= 4, Z = 2 , e = 5),
    y = c(X = 0, Y = 0, Z = 2, e = 1) 
  )) %>% 
  tidy_dagitty()     # Add edge_type within pipe

# Plot DAG 
dag_conf_p <-  dag_confounder %>% 
  ggdag(., layout = "circle") + 
  theme_dag_blank(plot.caption = element_text(hjust = 1)) +
  # Nodes 
  geom_dag_node(aes(color = name)) +  
  # Labels dodged to avoid overlap
  geom_dag_label(aes(label = label), color = "#4c4c4c", nudge_x = 0.7, nudge_y = 0.2) +  
  geom_dag_text(color="white") +
  labs(title = "Causal map with CONFOUNDER (Z)" , #subtitle = " ",
        caption = ) +
  # Map colors to specific nodes
  scale_color_manual(values = c("X" = "#9b6723", "Y" = "#72aed8", "Z" = "#d02e4c",
                                "e" = "#A6A6A6"), 
                     guide = "none")   

dag_conf_p

library(dagitty) # Graphical Analysis of Structural Causal Models # Graphical Analysis of Structural Causal Models # Graphical Analysis of Structural Causal Models 
library(ggdag) # Analyze and Create Elegant Directed Acyclic Graphs # Analyze and Create Elegant Directed Acyclic Graphs # Analyze and Create Elegant Directed Acyclic Graphs 

# df DAG 
dag_med <- dagify(
  Y ~ Z + e,
  Z ~ X,
  
  exposure = "X",
  outcome = "Y",
  #latent = "e",
  # labels
  labels = c(
    "Z" = "Mediator",
    "X" = "Exposure",
    "Y" = "Outcome",
    "e" = "Unobserved \nerror"),
  # positions
  coords = list(
    x = c(X = 0, Y= 4, Z = 2 , e = 5),
    y = c(X = 0, Y = 0, Z = 2, e = 1) 
  )) %>% 
  tidy_dagitty()     # Add edge_type within pipe

# Plot DAG 
dag_med_p <-  dag_med %>% 
  ggdag(., layout = "circle") +  # decided in dagify
  theme_dag_blank(plot.caption = element_text(hjust = 1)) +
  # Nodes 
  geom_dag_node(aes(color = name)) +  
  # Labels dodged to avoid overlap
  geom_dag_label(aes(label = label), color = "#4c4c4c", nudge_x = 0.7, nudge_y = 0.2) +  
  
  geom_dag_text(color="white") +
  labs(title = "Causal map with MEDIATOR (Z)" , 
       #subtitle = "X = exposure, Y = outcome, Z = collider, e = random error", 
       caption = ) +
  # Map colors to specific nodes
  scale_color_manual(values = c("X" = "#9b6723", "Y" = "#72aed8", "Z" = "#d02e4c",
                                "e" = "#A6A6A6"), 
                     guide = "none")    # Remove legend

dag_med_p

data(NHANES)
# Select relevant columns and drop rows with missing values
nhanes_conf <- NHANES %>%
  select(ID, Age, SmokeNow, BPSysAve) %>%
  drop_na()
  
# Take a quick look at the data
paint(nhanes_conf)

# linear rel 
qplot (Age, BPSysAve, data = nhanes_conf, color = SmokeNow) +
  geom_smooth(method = "lm", se = FALSE) +
  labs(title = "Scatterplot of Age and Blood Pressure by Smoking Status",
       x = "Age",
       y = "Blood Pressure (mmHg)",
       color = "Smoking Status")

# Unadjusted linear model 
model_unadjusted <- lm(BPSysAve ~ SmokeNow, data = nhanes_conf)
summary(model_unadjusted)

# Adjusted model
model_adjusted <- lm(BPSysAve ~ SmokeNow + Age, data = nhanes_conf)
summary(model_adjusted)

conf_sum <- modelsummary::msummary(list( "NO Confounder" = model_unadjusted, 
                                         "Confounder" = model_adjusted), 
                                   fmt="%.3f", # format
                                   gof_omit = 'DF|Deviance|Log.Lik.|F|R2 Adj.|AIC|BIC',
                                   stars=c('*' = .05, '**' = .01)#,
                                   #coef_rename = cm,
                                   #output = paste(output,"ch19_reg-R.tex",sep="")
) 

conf_sum

# Add predicted values from the unadjusted and adjusted models
nhanes_conf <- nhanes_conf %>%
  mutate(
    pred_unadjusted = predict(model_unadjusted),  # Predicted BPSysAve from unadjusted model
    pred_adjusted = predict(model_adjusted)       # Predicted BPSysAve from adjusted model
  )

# Reshape the data into long format using pivot_longer()
nhanes_long <- nhanes_conf %>%
  pivot_longer(cols = c(pred_unadjusted, pred_adjusted), 
               names_to = "model", values_to = "predicted") %>%
  mutate(model = recode(model, 
                        pred_unadjusted = "Unadjusted: SmokeNow on BPSysAve", 
                        pred_adjusted = "Adjusted: SmokeNow + Age on BPSysAve"))

# Violin plot with dashed line connecting the two conditions
ggplot(nhanes_long, aes(x = factor(SmokeNow), y = BPSysAve, fill = model)) +
  
  # Create the violin plot to show the distribution of blood pressure
  geom_point(color = "grey", alpha = 0.4, position = position_dodge(width = 0.3) ) +
  
  # Overlay the predicted dashed line between the two smoking conditions
  geom_line(aes(y = predicted, group = model, color = model), 
            size = 0.8, linetype = "dashed", position = position_dodge(width = 0.3)) +
  
  # Facet vertically for unadjusted and adjusted models
  facet_wrap(model ~ ., scales = "free_y",ncol = 2) +
  
  # Add labels and title
  labs(title = "Confounder Analysis: Predicted values in Unadjusted vs Adjusted Regression models",
       subtitle = "Age = Confounder",
       x = "Smoking Status",
       y = "Systolic Blood Pressure (BPSysAve)",
       color = "Model",
       fill = "Model"
       ) +
  
  # Customize colors for better contrast
  scale_fill_manual(values = c("Unadjusted: SmokeNow on BPSysAve" ="lightcoral" , 
                               "Adjusted: SmokeNow + Age on BPSysAve" = "lightblue")) +
  scale_color_manual(values = c("Unadjusted: SmokeNow on BPSysAve" = "red", 
                                "Adjusted: SmokeNow + Age on BPSysAve" = "blue")) +
  
  # Minimal theme for clarity
  theme_minimal()+
  theme(legend.position = "none")

data(NHANES)
# Select relevant columns and drop rows with missing values
nhanes_med <- NHANES %>%
  select(Gender, Age, BPSysAve, BMI, PhysActiveDays, Diabetes, SmokeNow) %>%
  drop_na()

# Take a quick look at the data
summary(nhanes_med)

# Unadjusted model: total effect of smoking on blood pressure
total_effect_model <- lm(BPSysAve ~ PhysActiveDays, data = nhanes_med)
summary(total_effect_model)

# Mediator model: SmokeNow affects BMI
mediator_model <- lm(BMI ~ PhysActiveDays, data = nhanes_med)
summary(mediator_model)

# Outcome model: PhysActiveDays and BMI affect BPSysAve
outcome_model <- lm(BPSysAve ~  PhysActiveDays + BMI, data = nhanes_med)
summary(outcome_model)

med_sum <- modelsummary::msummary(list("BMI ~ Mediator Effect" = mediator_model,
                                       "BPSysAve ~ Total Effect" = total_effect_model,
                                       "BPSysAve ~ Outcome" = outcome_model), 
                                  fmt="%.3f", # format
                                  gof_omit = 'DF|Deviance|Log.Lik.|F|R2 Adj.|AIC|BIC',
                                  stars=c('*' = .05, '**' = .01)#,
                                  #coef_rename = cm,
                                  #output = paste(output,"ch19_reg-R.tex",sep="")
) 

med_sum

# Check the names of the coefficients in models...
names(coef(mediator_model))     # [1] BMI ~ PhysActiveDays 

names(coef(outcome_model))      # [1] BPSysAve ~  PhysActiveDays + BMI    

names(coef(total_effect_model)) # [1] BPSysAve ~ PhysActiveDays

# BREAKING DOWN THE ESTIMATE OF PhysActiveDays COEFFICIENT 

# TOTAL effect: The effect of PhysActiveDays on BPSysAve without adjusting for BMI
total_effect <- coef(total_effect_model)["PhysActiveDays"]
total_effect # 0.59

# DIRECT effect: The effect of PhysActiveDays on BPSysAve after adjusting for BMI.
direct_effect <- coef(outcome_model)["PhysActiveDays"]  # Direct effect (controlling for BMI)
direct_effect # 0.621 

# Indirect effect: The portion of the effect on BPSysAve that is mediated through BMI
indirect_effect <- coef(mediator_model)["PhysActiveDays"] * coef(outcome_model)["BMI"]   
indirect_effect # [-0.0821 X 0.382] = -0.0313 

# or 
total_effect - direct_effect # !!!!!!!!  -0.0313

# Proportion mediated: The proportion of the total effect that is mediated through BMI
proportion_mediated <- glue::glue("{round(indirect_effect / total_effect *100, 2)}%")
proportion_mediated #  -0.0531

# Create a data frame for the effects
effects_df <- data.frame(
  Effect = c("Total Effect", "Direct Effect", "Indirect Effect"),
  Value = c(total_effect, direct_effect, indirect_effect)
)

# Load ggplot2
library(ggplot2)

# Create the bar plot
ggplot(effects_df, aes(x = Effect, y = Value, fill = Effect)) +
  geom_bar(stat = "identity") +
  labs(title = "Mediation Analysis: Total, Direct, and Indirect Effects", 
       subtitle = "Effect (coefficient) of Physical Activity on Systolic Blood Pressure (BPSysAve)",
       y = "", x = "") +
  theme_minimal()

library(mediation)
# recall
mediator_model[["call"]][["formula"]]

outcome_model[["call"]][["formula"]]

# Calculate the mediation effect
mediation_model <- mediate(mediator_model, outcome_model, 
                           treat = "PhysActiveDays", 
                           mediator = "BMI", 
                           boot=TRUE, sims=500)
summary(mediation_model)

#                Estimate 95% CI Lower 95% CI Upper p-value  
# ACME            -0.0313      -0.1024         0.04   0.324     = indirect or mediation effect
# ADE              0.6213       0.1479         1.16   0.024 *   = direct effect 
# Total Effect     0.5899       0.1354         1.15   0.020 *   = total effect 
# Prop. Mediated  -0.0531      -0.2969         0.10   0.344     = proportion mediated

# Add predicted values for BMI, BPSysAve (adjusted), and BPSysAve (total effect)
nhanes_med_pred <- nhanes_med %>%
  mutate(
    pred_bmi = predict(mediator_model),         # Predicted BMI (mediator model)
    pred_bps_adj = predict(outcome_model),      # Predicted BPSysAve (adjusted for Age)
    pred_bps_total = predict(total_effect_model) # Predicted BPSysAve (total effect)
  )

# Check column names
colnames(nhanes_med_pred)

# Reshape the data into long format for faceting
nhanes_long <- nhanes_med_pred %>%
  gather(key = "model", value = "predicted",  pred_bps_adj, pred_bps_total) %>%
  mutate(model = recode(model, 
                        #pred_bmi = "Mediator Effect: BPSysAve on BMI", 
                        pred_bps_adj = "Adjusted Effect: PhysActiveDays + BMI on BPSysAve", 
                        pred_bps_total = "Total Effect: PhysActiveDays on BPSysAve"))

# Plot with vertically aligned facets
ggplot(nhanes_long, aes(x = BMI)) +
  geom_point(aes(y = BPSysAve), color = "black", shape = 16, alpha = 0.5) +  # Actual data points for BPSysAve
  geom_line(aes(y = predicted, color = model), size = 1) +  # Predicted values from different models
  facet_grid(model ~ ., scales = "free_y") +  # Facet vertically
  labs(title = "Mediation Analysis: Total, Adjusted, and Mediator Effects",
       x = "BMI",
       y = "Outcome / Mediator",
       color = "Effect Type") +
  theme_minimal()

# Select relevant columns and drop rows with missing values
dat_conf <- dat %>%
  select(id, age, bp_systolic, fv, exerc, bmi, smoker, gender, hh_income_usd, ln_hh_income_percap ) %>%
  drop_na() %>% 
  filter(fv <=1500)

# Take a quick look at the data
summary(dat_conf)

Hmisc::describe(dat_conf$bp_systolic) # Y
Hmisc::describe(dat_conf$bmi) # X
Hmisc::describe(dat_conf$age) # Z (confounder )
Hmisc::describe(dat_conf$exerc)  # Z (confounder )
Hmisc::describe(dat_conf$ln_hh_income_percap) # Z (confounder )

# Unadjusted model
model_unadjusted <- lm( bp_systolic ~ bmi + ln_hh_income_percap , data = dat_conf)
summary(model_unadjusted)

# Adjusted model
model_adjusted <- lm(bp_systolic ~ bmi + age + ln_hh_income_percap, data = dat_conf)
summary(model_adjusted)

model_unadjusted$coefficients 
model_adjusted$coefficients

# Get predicted values using augment and attach to the original dataset
dat_unadjusted <- broom::augment(model_unadjusted, data = dat_conf) %>%
  select( -.resid,  -.hat, -.sigma, -.cooksd, -.std.resid ) %>%
  rename(pred_unadjusted = .fitted)

dat_adjusted <- broom::augment(model_adjusted, data = dat_conf) %>%
  select( -.resid,     -.hat, -.sigma, -.cooksd, -.std.resid ) %>%
  rename(pred_adjusted = .fitted)

# Combine the predicted values from both models
dat_conf_long <- dat_unadjusted %>%
  left_join(dat_adjusted, by = "id", suffix=c("",".y"))%>%
  select(-ends_with(".y")) %>% 
  # Convert the data from wide to long format
  pivot_longer(cols = c(pred_unadjusted, pred_adjusted), 
               names_to = "model", 
               values_to = "predicted")

# Check the structure of the long dataset
str(dat_conf_long)

# check 
dat_conf_long %>% group_by(model) %>% 
  summarise(mean(predicted), sd(predicted))

ggplot(dat_conf, aes(x = bmi, y = bp_systolic)) +
  geom_point(color = "grey", alpha = 1) +  
  geom_point(data = dat_unadjusted, aes(x = bmi, y = pred_unadjusted), 
             color = "#e60066", alpha = 0.20) +
  geom_point(data = dat_adjusted, aes(x = bmi, y = pred_adjusted), 
             color = "#399B23", alpha = 0.20) +
  # Add labels and title
  labs(title = "Confounder Analysis: Unadjusted (pink) vs Adjusted (green) fitted values",
       x = "BMI",
       y = "Systolic Blood Pressure (BPSysAve)"
       )  +
    # Minimal theme for clarity
  theme_minimal() +
  theme(legend.position = "none")

# Unadjusted model
model_unadjusted <- lm( bp_systolic ~ bmi  , data = dat_conf)  
 # Adjusted model
model_adjusted <- lm(bp_systolic ~ bmi + age  , data = dat_conf)  

tidy(model_unadjusted, conf.int = TRUE ) %>% 
  ggplot(data = ., 
         mapping = aes(x = estimate, y = term, xmin = conf.low, xmax = conf.high)) +
  geom_pointrange() +
  labs(title = "Model Estimates of bmi on bp_systolic",
       x = "Coefficient Estimate",
       y = "Predictor",
       caption = "Models fit with OLS. Error bars show the 95% confidence interval.") +
  scale_y_discrete(labels = c("Intercept",  "bmi")) +
  ggpubr::theme_pubclean(flip = TRUE)

tidy(model_adjusted, conf.int = TRUE ) %>% 
  ggplot(data = ., 
         mapping = aes(x = estimate, y = term, xmin = conf.low, xmax = conf.high)) +
  geom_pointrange() +
  labs(title = "Model Estimates of bmi + age on bp_systolic",
       x = "Coefficient Estimate",
       y = "Predictor",
       caption = "Models fit with OLS. Error bars show the 95% confidence interval.") +
  scale_y_discrete(labels = c("Intercept", "age", "bmi")) +
  ggpubr::theme_pubclean(flip = TRUE)

fit_unadj_results <- tidy(model_unadjusted, conf.int = TRUE) %>% 
  mutate(model = "Model Unadjusted")

fit_adj_results <- tidy(model_adjusted, conf.int = TRUE) %>% 
  mutate(model = "Model Adjusted")

model_results <- bind_rows(fit_unadj_results, fit_adj_results)

ggplot(data = model_results, 
       aes(x = estimate, y = term, xmin = conf.low, xmax = conf.high, 
           color = model, shape = model)) +
  geom_pointrange(position = position_dodge(width = 0.5)) +
  labs(title = "Model Estimates of BMI and age on REM Sleep",
       x = "Coefficient Estimate",
       y = "Predictor",
       caption = "Models fit with OLS. Error bars show the 95% confidence interval.",
        color = "Model:",
       shape = "Model:") +
  scale_y_discrete(labels = c("Intercept", "age", "bmi")) +
  ggpubr::theme_pubclean(flip = FALSE)